import getopt
import sys
from datetime import datetime, timedelta

import numpy as np
import pandas as pd
from elasticsearch import Elasticsearch

from air_web.config.config import config
from air_web.data_platform import init_db
from air_web.dw.data_mapping import DateType


class YmStat:
    def __init__(self, yesterday):
        self.sql_engine = init_db()
        self.es = Elasticsearch(config["ES_HOST"])

        self.month_start_date = (yesterday.replace(day=1)).strftime("%Y-%m-%d")
        self.month_end_date = (
            (yesterday.replace(day=1) + timedelta(days=31)).replace(day=1)
        ).strftime("%Y-%m-%d")
        self.year_start_date = (datetime(yesterday.year, 1, 1)).strftime(
            "%Y-%m-%d"
        )
        self.year_end_date = (datetime(yesterday.year + 1, 1, 1)).strftime(
            "%Y-%m-%d"
        )

    def save_data_to_table(self, res_df, save_table):
        """存储数据"""
        if res_df.empty:
            print(f"写入{save_table},结果数据为空")
            return
        res_df.replace({np.nan: None}, inplace=True)
        self.sql_engine.update_df_by_id(res_df, save_table)
        print(f"写入{save_table},数据条数:{len(res_df)}")

    def get_range_df(self, table_name, date_type, start_date, end_date):
        sql = """select * 
                 from {table_name}
                 where date_type={date_type}
                   and data_date >= '{start_date}'
                   and data_date < '{end_date}'
              """.format(
            table_name=table_name,
            date_type=date_type,
            start_date=start_date,
            end_date=end_date,
        )
        df = self.sql_engine.query(sql)
        print(
            f"从{table_name}表中查询{start_date}到{end_date},date_type={date_type},数据条数:{len(df)}"
        )
        return df

    def get_stat_df(self, df, stat_field, stat_type):
        ascending = True if stat_type == "min" else False
        stat_field_list = (
            ["p_kt", "p_kt_total"]
            if stat_field == "p_kt"
            else ["p_total", "p_total_kt"]
        )
        stat_field_list = [f"{stat_type}_{col}" for col in stat_field_list]
        df = df.sort_values(stat_field_list, ascending=ascending)
        df.drop_duplicates(self.GROUP_COLS, inplace=True)

        stat_cols = [
            f"{stat_type}_{col}" for col in self.STAT_COLS_MAP[stat_field]
        ]
        cols = self.DEFAULT_COLS + stat_cols
        df = df[cols]
        return df

    def agg_month_or_year(self, df):
        max_kt_df = self.get_stat_df(df, "p_kt", "max")
        max_total_df = self.get_stat_df(df, "p_total", "max")
        min_kt_df = self.get_stat_df(df, "p_kt", "min")
        min_total_df = self.get_stat_df(df, "p_total", "min")
        mean_kt_df = (
            df.groupby(self.GROUP_COLS)["avg_p_kt"].mean().reset_index()
        )
        mean_total_df = (
            df.groupby(self.GROUP_COLS)["avg_p_total"].mean().reset_index()
        )
        max_tmp_df = df.groupby(self.GROUP_COLS)["max_tmp"].max().reset_index()
        min_tmp_df = df.groupby(self.GROUP_COLS)["min_tmp"].min().reset_index()

        # 合并成一个新的 DataFrame
        result = pd.merge(max_kt_df, max_total_df, on=self.DEFAULT_COLS)
        result = pd.merge(result, min_kt_df, on=self.DEFAULT_COLS)
        result = pd.merge(result, min_total_df, on=self.DEFAULT_COLS)
        result = pd.merge(result, mean_kt_df, on=self.GROUP_COLS)
        result = pd.merge(result, mean_total_df, on=self.GROUP_COLS)
        result = pd.merge(result, max_tmp_df, on=self.GROUP_COLS)
        result = pd.merge(result, min_tmp_df, on=self.GROUP_COLS)

        return result

    def agg_org_type_year(self):
        print(f"进行年统计{self.year_start_date}到{self.year_end_date}")
        month_df = self.get_range_df(
            self.table_name,
            DateType.MONTH,
            self.year_start_date,
            self.year_end_date,
        )
        month_df["data_date"] = self.year_start_date  # 改为当年1月1日
        month_df["date_type"] = DateType.YEAR
        year_df = self.agg_month_or_year(month_df)
        self.save_data_to_table(year_df, self.table_name)

    def agg_org_type_month(self):
        print(f"进行月统计{self.month_start_date}到{self.month_end_date}")
        day_df = self.get_range_df(
            self.table_name,
            DateType.DAY,
            self.month_start_date,
            self.month_end_date,
        )
        day_df["data_date"] = self.month_start_date  # 改为当月1日
        day_df["date_type"] = DateType.MONTH
        month_df = self.agg_month_or_year(day_df)
        self.save_data_to_table(month_df, self.table_name)


class OrgTypeYmStat(YmStat):
    DEFAULT_COLS = [
        "org_no",
        "org_name",
        "ad_org_name",
        "p_org_no",
        "type_id",
        "type_code",
        "p_type_id",
        "date_type",
        "data_date",
    ]
    GROUP_COLS = ["org_no", "type_id"]
    STAT_COLS_MAP = {
        "p_kt": ["p_kt", "p_kt_total", "p_kt_time"],
        "p_total": ["p_total", "p_total_kt", "p_total_time"],
    }

    def __init__(self, yesterday):
        super().__init__(yesterday)
        self.yesterday = yesterday
        self.table_name = config.get(
            "ORGNO_TYPEID_YMD_STAT", "orgno_typeid_ymd_stat"
        )

    def get_quarter_date(self):
        current_month = self.yesterday.month

        if 1 <= current_month <= 2:  # 去年第四季度，起始月份为12月
            quarter_start_date = datetime(self.yesterday.year - 1, 12, 1)
            quarter_end_date = datetime(self.yesterday.year, 3, 1)
        elif 3 <= current_month <= 5:  # 第一季度，起始月份为3月
            quarter_start_date = datetime(self.yesterday.year, 3, 1)
            quarter_end_date = datetime(self.yesterday.year, 6, 1)
        elif 6 <= current_month <= 8:  # 第二季度，起始月份为6月
            quarter_start_date = datetime(self.yesterday.year, 6, 1)
            quarter_end_date = datetime(self.yesterday.year, 9, 1)
        elif 9 <= current_month <= 11:  # 第三季度，起始月份为9月
            quarter_start_date = datetime(self.yesterday.year, 9, 1)
            quarter_end_date = datetime(self.yesterday.year, 12, 1)
        else:  # 第四季度，起始月份为12月
            quarter_start_date = datetime(self.yesterday.year, 12, 1)
            quarter_end_date = datetime(self.yesterday.year + 1, 3, 1)

        quarter_start_date = quarter_start_date.strftime("%Y-%m-%d")
        quarter_end_date = quarter_end_date.strftime("%Y-%m-%d")
        return quarter_start_date, quarter_end_date

    def agg_org_type_quarter(self):
        quarter_start_date, quarter_end_date = self.get_quarter_date()
        print(f"进行季度统计{quarter_start_date}到{quarter_end_date}")
        day_df = self.get_range_df(
            self.table_name,
            DateType.DAY,
            self.month_start_date,
            self.month_end_date,
        )
        day_df["data_date"] = quarter_start_date  # 改为当季1日
        day_df["date_type"] = DateType.QUARTER
        quarter_df = self.agg_month_or_year(day_df)
        self.save_data_to_table(quarter_df, self.table_name)

    def main(self):
        print("开始统计地区-行业")
        self.agg_org_type_month()
        self.agg_org_type_year()
        self.agg_org_type_quarter()


class ConsYmStat(YmStat):
    DEFAULT_COLS = [
        "cons_no",
        "on5",
        "on7",
        "date_type",
        "data_date",
    ]
    GROUP_COLS = ["cons_no"]
    STAT_COLS_MAP = {
        "p_kt": ["p_kt", "p_kt_total", "p_kt_time"],
        "p_total": ["p_total", "p_total_kt", "p_total_time"],
    }

    def __init__(self, yesterday):
        super().__init__(yesterday)
        self.table_name = config.get("CONS_YMD_STAT", "cons_ymd_stat")

    def main(self):
        print("开始统计用户")
        self.agg_org_type_month()
        self.agg_org_type_year()


if __name__ == "__main__":
    start_date = datetime.now() - timedelta(days=1)
    end_date = datetime.now()

    opts, args = getopt.getopt(sys.argv[1:], "s:e:")
    for opt, val in opts:
        if opt == "-s":
            start_date = val
        elif opt == "-e":
            end_date = val

    date_list = pd.date_range(
        start=start_date, end=end_date, freq="1 D"
    ).strftime("%Y-%m-%d")[:-1]
    step_dict = {"save_table": "area_cons_num"}
    for data_time in date_list:
        data_time = datetime.strptime(data_time, "%Y-%m-%d")
        OrgTypeYmStat(data_time).main()
        ConsYmStat(data_time).main()
